package com.ruoyi.web.controller.chat.util;

import com.lzhpo.chatgpt.entity.chat.ChatCompletionMessage;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;

import javax.validation.constraints.NotEmpty;
import java.util.List;

/**
 * TODO
 *
 * @author chenchuan2
 * @date 2023-05-07 23:17:14
 */
public class ChatUtil {

    private static String DEFAULT_ROLE = "user";
    private static String SYSTEM_ROLE = "system";

    public static String getUserQa(@NotEmpty List<ChatCompletionMessage> message) {
        for (int i = message.size() - 1; i >= 0; i--) {
            ChatCompletionMessage chatCompletionMessage = message.get(i);
            if (StringUtils.endsWithIgnoreCase(DEFAULT_ROLE, chatCompletionMessage.getRole())) {
                return chatCompletionMessage.getContent();
            }
        }
        return "";
    }

    public static ChatCompletionMessage getSystem(@NotEmpty List<ChatCompletionMessage> message) {
        for (int i = 0; i < message.size(); i++) {
            ChatCompletionMessage chatCompletionMessage = message.get(i);
            if (StringUtils.endsWithIgnoreCase(SYSTEM_ROLE, chatCompletionMessage.getRole())) {
                return chatCompletionMessage;
            }
        }
        return null;
    }

    public static void appendSystem(@NotEmpty List<ChatCompletionMessage> message, String append) {
        ChatCompletionMessage system = getSystem(message);
        if (null == system) {
            system = new ChatCompletionMessage();
            system.setRole(SYSTEM_ROLE);
            system.setContent(append);
            message.add(0, system);
        } else {
            system.setContent(system.getContent() + ";" + append);
        }
    }
}
